package org.example.algorithm.tree;

//树中距离之和

import java.util.*;

/**
输入: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释: 树如图所示。
我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。
 */
public class SumOfDistancesInTreeSolution {
    private List<Integer>[] neighborInfo;
    private int[] ans;
    private int[] size;

    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        neighborInfo = new ArrayList[n];
        Arrays.setAll(neighborInfo, e->new ArrayList<>());
        for (int[] e: edges) {
            int x = e[0];
            int y = e[1];
            neighborInfo[x].add(y);
            neighborInfo[y].add(x);
        }
        ans = new int[n];
        size = new int[n];
        dfs(0, -1, 0);
        dfs2(0, -1);
        return ans;
    }

    private void dfs2(int x, int parent) {
        for (Integer node: neighborInfo[x]) {
            if (node != parent) {
                ans[node] = ans[x] + neighborInfo.length - 2 * size[node];
                dfs2(node, x);
            }
        }
    }

    //depth代表x到根节点距离
    public void dfs(int x, int parent, int depth) {
        ans[0] += depth;
        //初始化为1
        size[x] = 1;
        for (Integer node: neighborInfo[x]) {
            if (node != parent) {
                dfs(node, x, depth+1);
                size[x] += size[node];
            }
        }
    }

    public static void main(String[] args) {

    }
}
